import torch
import queue,os
import  threading  
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
from datasets import load_dataset
# from audio_tools import record_audio
from asa import record_audio


class WhisperModel(object):
    def __init__(self) -> None:
        device = "cuda:0" if torch.cuda.is_available() else "cpu"
        torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
        model_id = "/data/models/huggingface/whisper-large-v3"
        model = AutoModelForSpeechSeq2Seq.from_pretrained(
            model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True
        )
        model.to(device)
        processor = AutoProcessor.from_pretrained(model_id)
        self.SpeechToEngPipe = pipeline(
            "automatic-speech-recognition",
            model=model,
            tokenizer=processor.tokenizer,
            feature_extractor=processor.feature_extractor,
            torch_dtype=torch_dtype,
            device=device,
        )
        self.TranslatePipe = pipeline("translation_en_to_zh", model="/data/models/huggingface/opus-mt-en-zh",device=0)
        self.q = queue.Queue()

    def speech2chinese(self,eng_speech):
        result = self.SpeechToEngPipe(eng_speech)
        ch_result = self.TranslatePipe(result["text"])[0]["translation_text"]
        return ch_result

    def translate(self):
        i = 0
        while True:
            file = "output{}.mp3".format(i)
            record_audio(file)
            self.q.put(file)
            i+=1    
    
    def translate_thread(self):
        while True:
            file = self.q.get()
            print(self.speech2chinese(file))
            os.remove(file)
            print("######################################################################################")

    def translate_to_mp3(self,file):
        print(self.speech2chinese(file))

    def speech_to_english(self,file):
        print(self.SpeechToEngPipe(file))

wm = WhisperModel()
t = threading.Thread(target=wm.translate)  
t.start()  
wm.translate_thread()

# wm.translate_to_mp3("20240818-195501.mp3")
# wm.speech_to_english("20240818-195501.mp3")